#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
可视化模块
用于绘制股票价格、技术指标和交易信号
"""

import matplotlib
# 设置matplotlib后端，避免在PyCharm中出现兼容性问题
matplotlib.use('TkAgg')

import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from matplotlib.pylab import mpl
import numpy as np


# 设置中文字体，解决中文显示问题
# 检查系统中可用的中文字体
import matplotlib.font_manager as fm

# 尝试查找可用的中文字体
chinese_fonts = [f.name for f in fm.fontManager.ttflist if 'Sim' in f.name or 'Kai' in f.name or 'Hei' in f.name]

if chinese_fonts:
    # 如果找到中文字体，使用第一个
    mpl.rcParams['font.sans-serif'] = [chinese_fonts[0]]
else:
    # 如果没有找到中文字体，使用系统默认字体并设置支持中文
    mpl.rcParams['font.sans-serif'] = ['DejaVu Sans', 'Bitstream Vera Sans', 'sans-serif']
    
mpl.rcParams['axes.unicode_minus'] = False    # 解决保存图像是负号'-'显示为方块的问题


def plot_sh_index_data(data, signals=None, title="上证指数"):
    """
    绘制上证指数价格图
    
    参数:
    data: 股票价格数据
    signals: 交易信号数据（可选）
    title: 图表标题
    """
    fig, ax = plt.subplots(figsize=(12, 6))
    
    # 绘制收盘价
    ax.plot(data['date'], data['close'], label='收盘价', color='black', linewidth=1)
    
    # 如果有信号数据，则绘制买卖信号
    if signals is not None:
        # 绘制买入信号
        buy_signals = signals[signals['positions'] == 1.0]
        ax.plot(buy_signals['date'], buy_signals['close'], '^', markersize=10, color='green', label='买入信号')
        
        # 绘制卖出信号
        sell_signals = signals[signals['positions'] == -1.0]
        ax.plot(sell_signals['date'], sell_signals['close'], 'v', markersize=10, color='red', label='卖出信号')
        
        # 绘制移动平均线
        ax.plot(signals['date'], signals['short_ma'], label=f'短期均线', alpha=0.7)
        ax.plot(signals['date'], signals['long_ma'], label=f'长期均线', alpha=0.7)
    
    # 设置图表标题和标签
    ax.set_title(title)
    ax.set_xlabel('日期')
    ax.set_ylabel('价格')
    ax.legend()
    
    # 自动格式化日期显示
    ax.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m'))
    ax.xaxis.set_major_locator(mdates.MonthLocator(interval=3))
    plt.setp(ax.xaxis.get_majorticklabels(), rotation=45)
    
    plt.tight_layout()
    plt.show()


def plot_portfolio_value(portfolio_data, title="投资组合价值"):
    """
    绘制投资组合价值变化图
    
    参数:
    portfolio_data: 投资组合数据
    title: 图表标题
    """
    fig, ax = plt.subplots(figsize=(12, 6))
    
    # 绘制总资产价值
    ax.plot(portfolio_data['date'], portfolio_data['total_asset'], label='总资产', color='blue', linewidth=1)
    
    # 设置图表标题和标签
    ax.set_title(title)
    ax.set_xlabel('日期')
    ax.set_ylabel('资产价值')
    ax.legend()
    
    # 自动格式化日期显示
    ax.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m'))
    ax.xaxis.set_major_locator(mdates.MonthLocator(interval=3))
    plt.setp(ax.xaxis.get_majorticklabels(), rotation=45)
    
    plt.tight_layout()
    plt.show()